import pandas as pd
import numpy as np
import os
from typing import List, Dict
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class PredictionFormatter:
    """预测结果格式化器：为测试集预留预测字段"""

    def __init__(self, data_dir: str = None):
        self.data_dir = data_dir or r"D:\sjysds\pythonProject1\data"

    def get_data_path(self, relative_path: str) -> str:
        """获取数据目录下的绝对路径"""
        if os.path.isabs(relative_path):
            return relative_path
        return os.path.join(self.data_dir, relative_path)

    def add_prediction_fields(self, test_file: str = "test_processed.csv",
                              output_file: str = "test_with_prediction_fields.csv") -> pd.DataFrame:
        """
        为测试集添加预测字段

        Args:
            test_file: 处理后的测试集文件名
            output_file: 输出文件名

        Returns:
            添加预测字段后的测试集DataFrame
        """
        test_path = self.get_data_path(test_file)

        if not os.path.exists(test_path):
            raise FileNotFoundError(f"测试集文件未找到: {test_path}")

        # 读取测试集
        df = pd.read_csv(test_path, encoding='GBK')
        logger.info(f"加载测试集: {len(df)} 条")

        # 添加预测字段（使用正确的数据类型）
        df = df.copy()
        df['predicted_labels'] = None  # 初始为None，后续会转换为字符串
        df['predicted_options'] = ''  # 字符串类型

        # 确保test_question_id字段存在
        if 'test_question_id' not in df.columns:
            if 'index' in df.columns:
                df['test_question_id'] = df['index']
            else:
                df['test_question_id'] = range(1, len(df) + 1)

        # 显式设置数据类型
        df['predicted_labels'] = df['predicted_labels'].astype('object')  # 对象类型可以存储任何数据
        df['predicted_options'] = df['predicted_options'].astype(str)  # 明确设置为字符串类型

        # 保存结果
        output_path = self.get_data_path(output_file)
        df.to_csv(output_path, encoding='GBK', index=False)
        logger.info(f"已添加预测字段，保存至: {output_path}")

        # 显示字段信息
        logger.info("新增字段:")
        logger.info("  - predicted_labels: 二进制预测标签 (如 [1,0,1,0,0])")
        logger.info("  - predicted_options: 选项字符串 (如 'AC')")
        logger.info("  - test_question_id: 题目标识 (保留字段)")

        return df

    @staticmethod
    def labels_to_options(labels, threshold: float = 0.5) -> str:
        """
        将二进制标签转为选项字符串

        Args:
            labels: 二进制标签列表或数组 (如 [1,0,1,0,0])
            threshold: 阈值，大于等于此值视为选中

        Returns:
            选项字符串 (如 "AC")
        """
        if labels is None or (isinstance(labels, (list, np.ndarray)) and len(labels) == 0):
            return ""

        try:
            # 处理不同的输入格式
            if isinstance(labels, str):
                # 尝试解析字符串格式的列表
                if labels.startswith('[') and labels.endswith(']'):
                    labels = eval(labels)
                else:
                    return ""

            # 确保是列表或数组
            if not isinstance(labels, (list, np.ndarray)):
                return ""

            # 映射选项
            option_map = {0: 'A', 1: 'B', 2: 'C', 3: 'D', 4: 'E'}
            selected_options = []

            for i, val in enumerate(labels):
                if i < 5:  # 只处理前5个选项
                    # 处理不同的数值类型
                    if isinstance(val, (int, float, np.number)):
                        if float(val) >= threshold:
                            selected_options.append(option_map[i])
                    elif isinstance(val, str):
                        # 处理字符串类型的数值
                        try:
                            if float(val) >= threshold:
                                selected_options.append(option_map[i])
                        except ValueError:
                            continue

            return "".join(selected_options)

        except Exception as e:
            logger.error(f"标签转换失败: {e}")
            return ""

    def update_predictions(self, test_file: str = "test_with_prediction_fields.csv",
                           output_file: str = "test_with_predictions.csv"):
        """
        更新预测结果（模拟使用）

        Args:
            test_file: 带预测字段的测试集文件名
            output_file: 输出文件名
        """
        test_path = self.get_data_path(test_file)

        if not os.path.exists(test_path):
            raise FileNotFoundError(f"测试集文件未找到: {test_path}")

        # 读取测试集
        df = pd.read_csv(test_path, encoding='GBK')

        # 确保数据类型正确
        df['predicted_labels'] = df['predicted_labels'].astype('object')
        df['predicted_options'] = df['predicted_options'].astype(str)

        # 模拟一些预测结果（实际中由模型产生）
        np.random.seed(42)  # 为了可重现性

        predicted_labels_list = []
        predicted_options_list = []

        for idx in range(len(df)):
            # 生成随机的二进制标签（模拟模型预测）
            labels = [1 if np.random.random() > 0.7 else 0 for _ in range(5)]
            predicted_labels_list.append(str(labels))

            # 转换为选项字符串
            options_str = self.labels_to_options(labels)
            predicted_options_list.append(options_str)

        # 批量赋值，避免逐个赋值警告
        df['predicted_labels'] = predicted_labels_list
        df['predicted_options'] = predicted_options_list

        # 保存结果
        output_path = self.get_data_path(output_file)
        df.to_csv(output_path, encoding='GBK', index=False)
        logger.info(f"已更新预测结果，保存至: {output_path}")

        # 显示预测统计
        option_counts = df['predicted_options'].apply(len)
        logger.info(f"\n预测结果统计:")
        logger.info(f"平均选择选项数: {option_counts.mean():.2f}")
        logger.info(f"选项数分布:")
        for count in range(6):
            count_num = (option_counts == count).sum()
            percentage = count_num / len(df) * 100
            logger.info(f"  {count}个选项: {count_num}题 ({percentage:.1f}%)")

        return df

    def generate_submission(self, test_file: str = "test_with_predictions.csv",
                            submission_file: str = "submission.csv"):
        """
        生成最终提交文件

        Args:
            test_file: 带预测结果的测试集文件名
            submission_file: 提交文件名
        """
        test_path = self.get_data_path(test_file)

        if not os.path.exists(test_path):
            raise FileNotFoundError(f"测试集文件未找到: {test_path}")

        # 读取测试集
        df = pd.read_csv(test_path, encoding='GBK')

        # 生成提交文件
        submission = df[['test_question_id', 'predicted_options']].copy()
        submission = submission.rename(columns={'test_question_id': 'index'})

        # 确保predicted_options不为空
        submission['predicted_options'] = submission['predicted_options'].fillna('')

        # 保存提交文件
        submission_path = self.get_data_path(submission_file)
        submission.to_csv(submission_path, encoding='GBK', index=False)
        logger.info(f"提交文件已生成: {submission_path}")

        # 显示提交文件示例
        logger.info("\n提交文件示例:")
        print(submission.head(10).to_string(index=False))

        return submission

    def process_full_pipeline(self):
        """完整的处理流程"""
        try:
            # 1. 添加预测字段
            logger.info("步骤1: 添加预测字段")
            test_with_fields = self.add_prediction_fields()

            # 2. 模拟更新预测结果（实际中由模型完成）
            logger.info("步骤2: 模拟预测结果")
            test_with_predictions = self.update_predictions()

            # 3. 生成提交文件
            logger.info("步骤3: 生成提交文件")
            submission = self.generate_submission()

            logger.info("处理流程完成！")
            return submission

        except Exception as e:
            logger.error(f"处理流程失败: {e}")
            import traceback
            traceback.print_exc()


# 使用示例
def main():
    """主函数"""
    import argparse

    parser = argparse.ArgumentParser(description='预测结果格式化工具')
    parser.add_argument('--action', choices=['add_fields', 'simulate', 'submit', 'full'],
                        default='full', help='执行的操作')
    parser.add_argument('--input', default='test_processed.csv', help='输入文件名')
    parser.add_argument('--output', default='test_with_prediction_fields.csv', help='输出文件名')
    parser.add_argument('--submission', default='submission.csv', help='提交文件名')

    args = parser.parse_args()

    formatter = PredictionFormatter()

    try:
        if args.action == 'add_fields':
            # 只添加预测字段
            formatter.add_prediction_fields(args.input, args.output)

        elif args.action == 'simulate':
            # 模拟预测结果
            formatter.update_predictions(args.input, args.output)

        elif args.action == 'submit':
            # 生成提交文件
            formatter.generate_submission(args.input, args.submission)

        else:
            # 完整流程
            formatter.process_full_pipeline()

    except Exception as e:
        logger.error(f"处理失败: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()
